import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np
import cv2
import os


class MyDataset(Dataset):
    def __init__(self, mask_dir, image_dir, transform=None):
        self.mask_dir = mask_dir
        self.image_dir = image_dir
        self.transform = transform

        self.mask_list = []
        self.image_list = []
        self.mask_list = sorted(os.listdir(self.mask_dir))
        self.image_list = sorted(os.listdir(self.image_dir))

    def __len__(self):
        return len(self.mask_list)

    def __getitem__(self, idx):
        mask_path = os.path.join(self.mask_dir, self.mask_list[idx])
        image_path = os.path.join(self.image_dir, self.image_list[idx])

        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
        image = cv2.imread(image_path, cv2.IMREAD_COLOR)

        if self.transform:
            mask = self.transform(mask)
            image = self.transform(image)

        # 转换为PyTorch张量
        mask = torch.tensor(mask, dtype=torch.float32)
        image = torch.tensor(image, dtype=torch.float32)

        # 计算图像的傅里叶变换
        image_fft = np.fft.fft2(image.numpy())  # 转为NumPy进行FFT计算
        image_fft = np.fft.fftshift(image_fft)
        fft = np.log(np.abs(image_fft) + 1)

        # 转换FFT结果为PyTorch张量
        fft = torch.tensor(fft, dtype=torch.float32)

        return mask, image, fft


def get_dataloader(mask_dir, image_dir, batch_size=32, transform=None):
    dataset = MyDataset(mask_dir, image_dir, transform=transform)
    print("Dataset size: ", len(dataset))
    dataloader = DataLoader(
        dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn
    )
    return dataloader


def collate_fn(batch):
    # 找到每一维的最大大小
    max_h = max(img[1].shape[0] for img in batch)  # 找到最大高度
    max_w = max(img[1].shape[1] for img in batch)  # 找到最大宽度

    # 填充每个图像到最大大小
    padded_batch = []
    for mask, img, fft in batch:
        pad_h = max_h - img.shape[0]
        pad_w = max_w - img.shape[1]

        padded_img = torch.nn.functional.pad(
            torch.tensor(img), (0, 0, 0, pad_w, 0, pad_h)
        )  
        padded_mask = torch.nn.functional.pad(
            mask, (0, pad_w, 0, pad_h)
        )  
        padded_fft = torch.nn.functional.pad(
            torch.tensor(fft), (0, 0, 0, pad_w, 0, pad_h)
        )  
        padded_batch.append((padded_mask, padded_img, padded_fft))

    # 返回填充后的张量
    return [torch.stack([item[i] for item in padded_batch]) for i in range(3)]


# Example usage
if __name__ == "__main__":
    mask_dir = "/home/challenge/dataset/sd2sp/masks_expanded"
    image_dir = "/home/challenge/dataset/sd2sp/seg"
    dataloader = get_dataloader(mask_dir, image_dir, batch_size=16)

    for masks, images, images_fft in dataloader:
        print(masks.shape)
        print(images.shape)
        print(images_fft.shape)
